"""
任务2：KNN分类算法 - Iris数据集
从零实现K近邻算法，不使用scikit-learn
"""

import numpy as np
from collections import Counter
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split


class KNN:
    def __init__(self, k=3):
        self.k = k

    def fit(self, X, y):
        self.X_train = X
        self.y_train = y

    def predict(self, X):
        return np.array([self._predict(x) for x in X])

    def _predict(self, x):
        distances = [np.sqrt(np.sum((x - x_train) ** 2)) for x_train in self.X_train]
        k_indices = np.argsort(distances)[:self.k]
        k_labels = [self.y_train[i] for i in k_indices]
        return Counter(k_labels).most_common(1)[0][0]

    def accuracy(self, X_test, y_test):
        predictions = self.predict(X_test)
        return np.sum(predictions == y_test) / len(y_test)


def main():
    print("=== 任务2：KNN分类 - Iris数据集 ===\n")

    # 加载数据
    iris = load_iris()
    X, y = iris.data, iris.target

    print("数据集信息:")
    print(f"- 记录数: {X.shape[0]}")
    print(f"- 属性数: {X.shape[1]}")
    print(f"- 属性: {iris.feature_names}")
    print(f"- 类别: {iris.target_names}")

    # 划分训练测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

    # 训练KNN
    knn = KNN(k=3)
    knn.fit(X_train, y_train)
    accuracy = knn.accuracy(X_test, y_test)

    print(f"\n实验结果:")
    print(f"- KNN准确率: {accuracy:.4f} ({accuracy * 100:.2f}%)")

    # 不同K值比较
    print(f"\n不同K值准确率:")
    for k in [1, 3, 5, 7]:
        knn_temp = KNN(k=k)
        knn_temp.fit(X_train, y_train)
        acc = knn_temp.accuracy(X_test, y_test)
        print(f"- K={k}: {acc:.4f}")


if __name__ == "__main__":
    main()